package luke.audio.analyzer;

import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.*;

/**
 * 频谱分析器类
 * 提供音频信号的频谱分析功能，包括计算频谱图和频带能量
 */
public class SpectrumAnalyzer {
    // 采样率
    private final int sampleRate;
    // FFT大小
    private final int fftSize;
    // 快速傅里叶变换器
    private final FastFourierTransformer fft;

    /**
     * 构造函数
     * 
     * @param sampleRate 采样率
     * @param fftSize    FFT大小
     */
    public SpectrumAnalyzer(int sampleRate, int fftSize) {
        this.sampleRate = sampleRate;
        this.fftSize = fftSize;
        this.fft = new FastFourierTransformer(DftNormalization.STANDARD);
    }

    /**
     * 计算频谱图
     * 
     * @param audioData  音频数据
     * @param hopSize    步长
     * @return 频谱图二维数组
     */
    public double[][] computeSpectrogram(float[] audioData,  int hopSize) {
        //分帧处理
        int numWindows = (audioData.length - fftSize) / hopSize + 1;
        double[][] spectrogram = new double[numWindows][];

        // 应用汉宁窗
        double[] window = createHanningWindow(fftSize);

        for (int i = 0; i < numWindows; i++) {
            double[] windowedData = new double[fftSize];
            int start = i * hopSize;

            for (int j = 0; j < fftSize; j++) {
                windowedData[j] = audioData[start + j] * window[j];
            }

            // 执行FFT
            Complex[] fftResult = fft.transform(windowedData, TransformType.FORWARD);
            double[] magnitudeSpectrum = new double[fftSize / 2];

            for (int j = 0; j < fftSize / 2; j++) {
                magnitudeSpectrum[j] = fftResult[j].abs();
            }

            spectrogram[i] = magnitudeSpectrum;
        }

        return spectrogram;
    }

    /**
     * 计算单帧音频数据的频谱
     * 
     * @param audioData  音频数据（必须是fftSize长度）
     * @return 频谱数组
     */
    public double[] computeSpectrum(float[] audioData) {
        if (audioData.length != fftSize) {
            throw new IllegalArgumentException("音频数据长度必须等于FFT大小: " + fftSize);
        }

        // 应用汉宁窗
        double[] window = createHanningWindow(fftSize);
        double[] windowedData = new double[fftSize];

        for (int i = 0; i < fftSize; i++) {
            windowedData[i] = audioData[i] * window[i];
        }

        // 执行FFT
        Complex[] fftResult = fft.transform(windowedData, TransformType.FORWARD);
        double[] magnitudeSpectrum = new double[fftSize / 2];

        for (int i = 0; i < fftSize / 2; i++) {
            magnitudeSpectrum[i] = fftResult[i].abs();
        }

        return magnitudeSpectrum;
    }

    /**
     * 计算频带能量
     * 
     * @param spectrogram    频谱图
     * @param frequencyBands 频带边界数组
     * @return 频带能量二维数组
     */
    public double[][] computeBandEnergy(double[][] spectrogram, double[] frequencyBands) {
        // 频带数量应该是数组长度的一半，因为每两个元素定义一个频带的频率范围
        int numBands = frequencyBands.length / 2;
        int numWindows = spectrogram.length;
        double[][] bandEnergy = new double[numWindows][numBands];

        for (int i = 0; i < numWindows; i++) {
            for (int j = 0; j < numBands; j++) {
                // 每两个相邻的频率值定义一个频带，即frequencyBands[j*2]到frequencyBands[j*2+1]
                double lowFreq = frequencyBands[j * 2];
                double highFreq = frequencyBands[j * 2 + 1];
                bandEnergy[i][j] = computeBandEnergy(spectrogram[i], lowFreq, highFreq);
            }
        }

        return bandEnergy;
    }

    /**
     * 计算单个频带能量
     * 
     * @param spectrum  频谱数据
     * @param lowFreq   低频边界
     * @param highFreq  高频边界
     * @return 频带能量值
     */
    public double computeBandEnergy(double[] spectrum, double lowFreq, double highFreq) {
        // 使用更精确的频带边界计算，避免整数截断
        double lowBin = lowFreq * fftSize / sampleRate;
        double highBin = highFreq * fftSize / sampleRate;
        double energy = 0;

        // 处理边界部分的能量贡献
        int lowIndex = (int) Math.ceil(lowBin);
        int highIndex = (int) Math.floor(highBin);

        // 处理低频边界部分
        if (lowIndex > 0 && lowIndex < spectrum.length) {
            double lowFraction = lowIndex - lowBin;
            energy += spectrum[lowIndex - 1] * lowFraction;
        }

        // 处理高频边界部分
        if (highIndex >= 0 && highIndex < spectrum.length - 1) {
            double highFraction = highBin - highIndex;
            energy += spectrum[highIndex + 1] * highFraction;
        }

        // 处理完整频段的能量
        for (int i = Math.max(0, lowIndex); i <= Math.min(highIndex, spectrum.length - 1); i++) {
            energy += spectrum[i];
        }

        return energy;
    }

    /**
     * 创建汉宁窗
     * 
     * @param size 窗口大小
     * @return 汉宁窗数组
     */
    private double[] createHanningWindow(int size) {
        double[] window = new double[size];
        for (int i = 0; i < size; i++) {
            window[i] = 0.5 * (1 - Math.cos(2 * Math.PI * i / (size - 1)));
        }
        return window;
    }
}